import os
import time
import numpy as np
from sklearn.manifold import TSNE
import matplotlib.pyplot as plt
import torch
import torchvision as tv
import torchvision.transforms as transforms
from args import *
from tool import *
import evaluate
# from models import *


save_path = os.path.join(args.log_path, "image.{}".format(args.model_type))
if not os.path.exists(save_path):
    os.makedirs(save_path)


train_loader = torch.utils.data.DataLoader(
    tv.datasets.MNIST(args.data_path, train=True, download=True,
                   transform=transforms.Compose([
                    #    transforms.Resize([224, 224]),
                       transforms.ToTensor(),
                       transforms.Normalize((0.1307,), (0.3081,)),
                    #    transforms.Lambda(lambda x: x.repeat(3, 1, 1))
                   ])),
    batch_size=args.batch_size, shuffle=True)

test_loader = torch.utils.data.DataLoader(
    tv.datasets.MNIST(args.data_path, train=False,
                   transform=transforms.Compose([
                    #    transforms.Resize([224, 224]),
                       transforms.ToTensor(),
                       transforms.Normalize((0.1307,), (0.3081,)),
                    #    transforms.Lambda(lambda x: x.repeat(3, 1, 1))
                   ])),
    batch_size=100, shuffle=True)


P = os.path.join(args.log_path, "naive_0bias_0.0lambda_20scale_1dy-scale")
P_b = os.path.join(args.log_path, "naive_1bias_0.0lambda_20scale_1dy-scale")

model = torch.load(os.path.join(P, "26e_0.9869acc.pth")).cuda()
model_b = torch.load(os.path.join(P_b, "26e_0.9876acc.pth")).cuda()

C = np.load(os.path.join(P, "C.26e.npy"))
C_b = np.load(os.path.join(P_b, "C.26e.npy"))
N_CLASS = C.shape[0]


def get_fea():
    F, L, F_b = [], [], []
    _cnt = 0
    with torch.no_grad():
        for X, Y in test_loader:
            # X, Y = X.cuda(), Y.cuda()
            X = X.cuda()
            _f, _ = model(X)
            _fb, _ = model_b(X)
            F.append(_f.cpu().numpy())
            F_b.append(_fb.cpu().numpy())
            L.append(Y.numpy())
            _cnt += _f.size(0)
            if _cnt >= 2000:
                break

    F = np.vstack(F)
    F_b = np.vstack(F_b)
    L = np.concatenate(L)
    return F, F_b, L


def n_sample_pc():
    cnt = [0] * N_CLASS
    for _, Y in train_loader:
        for c in Y:
            cnt[c] += 1
    return cnt


F, F_b, L = get_fea()
W_t = model.state_dict()["clf_layer.weight"].cpu()
W = W_t  # [10, 2]
sample_pc = n_sample_pc()
norm_pc = W_t.norm(2, 1)
W_b = model_b.state_dict()["clf_layer.weight"].cpu().numpy()  # [10, 2]


title = "sample num and weight vector norm"
fig, ax = plt.subplots(2, 1)
plt.title(title)
ax[0].set_title("#samples per class")
ax[1].set_title("weight vectors' norm")
name = list(range(N_CLASS))
ax[0].bar(name, sample_pc, width=0.35)
ax[1].bar(name, norm_pc, width=0.35)
fig.savefig(os.path.join(save_path, "{}.png".format(title.replace(" ", "_"))))
plt.close(fig)

"""
title = "feature"
fig, ax = plt.subplots(1, 2)
plt.title(title)
ax[0].set_title("naive (no-bias)")
ax[1].set_title("naive")
ax[0].scatter(F[:, 0], F[:, 1], s=10, c=L, marker='.', cmap="rainbow")
ax[1].scatter(F_b[:, 0], F_b[:, 1], s=10, c=L, marker='.', cmap="rainbow")
fig.savefig(os.path.join(save_path, "{}.png".format(title)))
plt.close(fig)


title = "weight vectors"
fig, ax = plt.subplots(1, 2)
plt.title(title)
ax[0].set_title("naive (no-bias)")
ax[1].set_title("naive")
for i in range(W.shape[0]):
    ax[0].plot([0, W[i][0]], [0, W[i][1]])
    ax[1].plot([0, W_b[i][0]], [0, W_b[i][1]])
    ax[0].text(W[i][0], W[i][1], i)
    ax[1].text(W_b[i][0], W_b[i][1], i)
fig.savefig(os.path.join(save_path, "{}.png".format(title.replace(" ", "_"))))
plt.close(fig)


title = "centres and weight vectors"
fig, ax = plt.subplots(1, 2)
plt.title(title)
ax[0].set_title("naive (no-bias)")
ax[1].set_title("naive")
ax[0].scatter(C[:, 0], C[:, 1], s=50, c=np.arange(N_CLASS), marker='*', cmap="rainbow")
ax[1].scatter(C_b[:, 0], C_b[:, 1], s=50, c=np.arange(N_CLASS), marker='*', cmap="rainbow")
for i in range(W.shape[0]):
    # weight vector
    ax[0].plot([0, W[i][0]], [0, W[i][1]])
    ax[1].plot([0, W_b[i][0]], [0, W_b[i][1]])
    # class tag
    ax[0].text(C[i][0], C[i][1], i)
    ax[1].text(C_b[i][0], C_b[i][1], i)
    ax[0].text(W[i][0], W[i][1], i)
    ax[1].text(W_b[i][0], W_b[i][1], i)
fig.savefig(os.path.join(save_path, "{}.png".format(title.replace(" ", "_"))))
plt.close(fig)
"""
